import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.documents import Document
from langchain_community.vectorstores import Chroma

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPU_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")

embeddings  = ZhipuAIEmbeddings(api_key=api_key)

text = "你变廋了。"
text_list = [
    Document(page_content="你减肥成功了。",),
    Document(page_content="你变苗条了。",),
    Document(page_content="你又胖了。",),
    Document(page_content="你好廋啊！"),
]

db = Chroma.from_documents(text_list, embeddings)
retriever = db.as_retriever()

result = db.similarity_search(text)
print(f"\n===similarity_search===\n{result}")

result = db.similarity_search_with_score(text)
print(f"\n===similarity_search_with_score===\n{result}")

""" text_embedding = embeddings.embed_query(text)
result = db.similarity_search_by_vector(text_embedding)
print(f"\n===similarity_search_by_vector===\n{result}") """
